import numpy as np
from math import asinh, atan, sqrt, pi
#=======================================================================================
def f( p):
    eps = 1e-18
    x, y, z = abs(p[0]), abs(p[1]), abs(p[2])
    return + y / 2.0 * (z**2 - x**2) * asinh(y / (sqrt(x**2 + z**2) + eps)) \
         + z / 2.0 * (y**2 - x**2) * asinh(z / (sqrt(x**2 + y**2) + eps)) \
         - x*y*z * atan(y*z / (x * sqrt(x**2 + y**2 + z**2) + eps))       \
         + 1.0 / 6.0 * (2*x**2 - y**2 - z**2) * sqrt(x**2 + y**2 + z**2)

def g(p):
    eps = 1e-18
    x, y, z = p[0], p[1], abs(p[2])
    return + x*y*z * asinh(z / (sqrt(x**2 + y**2) + eps))                         \
         + y / 6.0 * (3.0 * z**2 - y**2) * asinh(x / (sqrt(y**2 + z**2) + eps)) \
         + x / 6.0 * (3.0 * z**2 - x**2) * asinh(y / (sqrt(x**2 + z**2) + eps)) \
         - z**3 / 6.0 * atan(x*y / (z * sqrt(x**2 + y**2 + z**2) + eps))        \
         - z * y**2 / 2.0 * atan(x*z / (y * sqrt(x**2 + y**2 + z**2) + eps))    \
         - z * x**2 / 2.0 * atan(y*z / (x * sqrt(x**2 + y**2 + z**2) + eps))    \
         - x*y * sqrt(x**2 + y**2 + z**2) / 3.0
def setNdemag(c, permute, n, dx, n_demag):
  print("正在计算退磁张量"+str(c))
  flags=np.array([#八部分各不相同
    [0,0,0, 0,0,0, 0,0],
    [0,1,1, 0,0,1, 1,0],
    [0,1,0, 1,1,0, 1,0],
    [0,0,0, 0,0,0, 0,0],
    [0,0,1, 1,1,1, 0,0],
    [0,0,0, 0,0,0, 0,0],
  ])
  flags=flags[c]
  func=[f,g,g,f,g,f][c]
  for index1 in range(n[0]):
    for index2 in range(n[1]):
      for index3 in range(n[2]):
          value = 0.0
          for i in np.rollaxis(np.indices((2,)*6), 0, 7).reshape(64, 6):
            value += (-1)**sum(i) * func(list(map(lambda j: ([index1,index2,index3][j] + i[j] - i[j+3]) * dx[j], permute)))
          n_demag[index1,index2,index3,c]=  value / (4 * pi * np.prod(dx))
  for index1 in range(1,n[0]):
    for index2 in range(n[1]):
      for index3 in range(n[2]):
          n_demag[-index1,index2,index3,c]= ((-1)**flags[1])*n_demag[index1,index2,index3,c]
  for index1 in range(n[0]):
    for index2 in range(1,n[1]):
      for index3 in range(n[2]):
          n_demag[index1,-index2,index3,c]=((-1)**flags[2])* n_demag[index1,index2,index3,c]
  for index1 in range(n[0]):
    for index2 in range(n[1]):
      for index3 in range(1,n[2]):
          n_demag[index1,index2,-index3,c]=((-1)**flags[3])* n_demag[index1,index2,index3,c]
  for index1 in range(1,n[0]):
    for index2 in range(1,n[1]):
      for index3 in range(n[2]):
          n_demag[-index1,-index2,index3,c]=((-1)**flags[4])* n_demag[index1,index2,index3,c]
  for index1 in range(1,n[0]):
    for index2 in range(n[1]):
      for index3 in range(1,n[2]):
          n_demag[-index1,index2,-index3,c]=((-1)**flags[5])* n_demag[index1,index2,index3,c]
  for index1 in range(n[0]):
    for index2 in range(1,n[1]):
      for index3 in range(1,n[2]):
          n_demag[index1,-index2,-index3,c]=((-1)**flags[6])* n_demag[index1,index2,index3,c]
  for index1 in range(1,n[0]):
    for index2 in range(1,n[1]):
      for index3 in range(1,n[2]):
          n_demag[-index1,-index2,-index3,c]=((-1)**flags[7])* n_demag[index1,index2,index3,c]
  return n_demag[:,:,:,c]